%%% rtmp_client.erl - barebones client for Flash rtmp
%%%
%%% The MIT License
%%%
%%% Copyright (c) 2008 hellomatty@gmail.com
%%%
%%% Permission is hereby granted, free of charge, to any person obtaining a copy
%%% of this software and associated documentation files (the "Software"), to deal
%%% in the Software without restriction, including without limitation the rights
%%% to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
%%% copies of the Software, and to permit persons to whom the Software is
%%% furnished to do so, subject to the following conditions:
%%%
%%% The above copyright notice and this permission notice shall be included in
%%% all copies or substantial portions of the Software.
%%%
%%% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
%%% IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
%%% FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
%%% AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
%%% LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
%%% OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
%%% THE SOFTWARE.
%%%

-module(rtmp_client).

-export([stream_audio/2]).
-include_lib("rtmp.hrl").
-include_lib("util.hrl").

-import(util).

-compile(export_all).

-record(rtmp_client_state, {
  socket,
  stream,
  stream_sink_pid = undefined,
  buffer_pid
}).

-record(rtmp_protocol_state,{
  socket,
  client_pid,
  amf_states = dict:new(),
  chunk_size = 128,
  bytes_read = 0
}).

-record(rtmp_amf_state, {
  header,
  chunks_length = 0,
  chunks = []
}).

-record(rtmp_header, {
  amf_index,
  type,
  body_length,
  stream
  }).

stream_audio({stream, Host, _StreamId, _Auth} = Stream, SinkPid) ->
  ?D({rtmp_client,stream,Stream}),
  spawn(fun() ->
    ?D({rtmp,client,self()}),
    link(SinkPid),
    {ok, Socket} = gen_tcp:connect(Host, ?RTMP_PORT, [binary, {packet, 0}, {active, false}]),
    ok = handshake(Socket),
    This = self(),
    BufferPid = spawn_link(fun() ->
      ?D({rtmp,buffer,self()}),
      consume_buffer(#rtmp_protocol_state{
        socket = Socket,
        client_pid = This
      })
    end),
    ok = invoke_stream(Socket, Stream),
    loop(#rtmp_client_state{socket = Socket, stream = Stream, stream_sink_pid = SinkPid, buffer_pid = BufferPid})
  end).

invoke_stream(Socket, Stream) ->
  ?D({connect}),
  send(Socket, ems_rtmp:encode( #channel{id=3, timestamp = 0, type = ?RTMP_TYPE_INVOKE, stream = 0}, amf_connect(Stream))),
  {ok, _} = wait_result(),
  %%send(Socket, rtmp_ping_clientbuffer(0,6000)),
  ?D({connect,ok}),

  ?D({createStream}),
  AMF_CALL_STREAM_CLIENT_NUMBER = 3,
  send(Socket, ems_rtmp:encode( #channel{id=3, timestamp = 0, type = ?RTMP_TYPE_INVOKE, stream = 0}, amf_create_stream(AMF_CALL_STREAM_CLIENT_NUMBER))),
  {ok, _} = wait_result(),
  ?D({createStream,ok}),

  ?D({play, Stream}),
  send(Socket, ems_rtmp:encode( #channel{id=3, timestamp = 0, type = ?RTMP_TYPE_INVOKE, stream = 1}, amf_play(Stream))),
  send(Socket, rtmp_ping_clientbuffer(1,300)),
  ok.

loop(ClientState) ->
  NewClientState = receive
    {rtmp,{rtmp_header,_,?RTMP_TYPE_NOTIFY,_,_},["onPlayStatus",{object,Status}]} ->
      case lists:member({code,"NetStream.Play.Complete"}, Status) of
        true ->
          exit(end_of_stream);
        _ ->
          ok
      end;
    {rtmp,{rtmp_header,_Index,?RTMP_TYPE_AUDIO,_,_}, <<_Tag:8,Frames/binary>>} ->
      case ClientState#rtmp_client_state.stream_sink_pid of
        Pid when is_pid(Pid) ->
          Pid ! {stream, from, self(), audio, Frames};
        _ ->
          ok  
      end,
      ClientState;
    {rtmp, _Header, _Body} = RTMPMsg ->
      ?D({proto, received, RTMPMsg}),
      ClientState;
    Backstop ->
      ?D({rtmp, got, Backstop}),
      ClientState
  end,
  loop(NewClientState).

consume_buffer(ProtocolState) ->
  Socket = ProtocolState#rtmp_protocol_state.socket,
  MetaHeader = get(Socket, 1),
  {HeaderLength, AMFIndex} = rtmp_decode_metaheader( MetaHeader ),
  HeaderBytes = case HeaderLength of
    1 ->
      MetaHeader;
    _ ->
      Bytes = get(Socket, HeaderLength - 1),
      list_to_binary([MetaHeader, Bytes])
  end,
  LastAMFState = try
    dict:fetch(AMFIndex, ProtocolState#rtmp_protocol_state.amf_states)
  catch
    error:_ -> #rtmp_amf_state{header = #rtmp_header{amf_index = AMFIndex}}
  end,
  LastHeader = LastAMFState#rtmp_amf_state.header,
  Header = rtmp_decode_header(HeaderBytes, LastHeader),
  Chunk = rtmp_get_chunk(Socket, Header#rtmp_header.body_length - LastAMFState#rtmp_amf_state.chunks_length, ProtocolState#rtmp_protocol_state.chunk_size),
  AMFState = LastAMFState#rtmp_amf_state{
    header = Header,
    chunks = [Chunk | LastAMFState#rtmp_amf_state.chunks],
    chunks_length = LastAMFState#rtmp_amf_state.chunks_length + size(Chunk)
  },
  NewProtocolState = process(ProtocolState, AMFState),
  consume_buffer(NewProtocolState).

process(ProtocolState, AMFState)
  when
    AMFState#rtmp_amf_state.chunks_length == (AMFState#rtmp_amf_state.header)#rtmp_header.body_length, 
    (AMFState#rtmp_amf_state.header)#rtmp_header.type == ?RTMP_TYPE_CHUNK_SIZE ->
  <<ChunkSize:32>> = list_to_binary( lists:reverse( AMFState#rtmp_amf_state.chunks ) ),
  ?D({rtmp,new_chunksize,ChunkSize}),
  ProtocolState#rtmp_protocol_state{
    amf_states = store_amf(ProtocolState, reset(AMFState)), 
    chunk_size = ChunkSize
  };
process(ProtocolState, AMFState)
  when
    AMFState#rtmp_amf_state.chunks_length == (AMFState#rtmp_amf_state.header)#rtmp_header.body_length, 
    (AMFState#rtmp_amf_state.header)#rtmp_header.type == ?RTMP_TYPE_BW_SERVER ->
  ?D({rtmp,server_bw,replying}),
  send(ProtocolState#rtmp_protocol_state.socket, rtmp_serverbw()),
  ProtocolState#rtmp_protocol_state{
    amf_states = store_amf(ProtocolState, reset(AMFState))
  };
process(ProtocolState, AMFState)
  when
    AMFState#rtmp_amf_state.chunks_length == (AMFState#rtmp_amf_state.header)#rtmp_header.body_length, 
    (AMFState#rtmp_amf_state.header)#rtmp_header.type == ?RTMP_TYPE_BW_CLIENT ->
  ?D({rtmp,client_bw,replying}),
  send(ProtocolState#rtmp_protocol_state.socket, rtmp_ping_clientbuffer(0,6000)),
  ProtocolState#rtmp_protocol_state{
    amf_states = store_amf(ProtocolState, reset(AMFState))
  };
process(ProtocolState, AMFState)
  when
    AMFState#rtmp_amf_state.chunks_length == (AMFState#rtmp_amf_state.header)#rtmp_header.body_length, 
    (AMFState#rtmp_amf_state.header)#rtmp_header.type == ?RTMP_TYPE_PING ->
  case list_to_binary( lists:reverse( AMFState#rtmp_amf_state.chunks ) ) of
    <<?PING_PING_CLIENT:16, Time:32>> ->
      ?D({rtmp,ping,pong}),
      send(ProtocolState#rtmp_protocol_state.socket, rtmp_ping_pong(Time));
    _ ->
      ok
  end,
  ProtocolState#rtmp_protocol_state{
    amf_states = store_amf(ProtocolState, reset(AMFState))
  };
process(ProtocolState, AMFState)
  when
    AMFState#rtmp_amf_state.chunks_length == (AMFState#rtmp_amf_state.header)#rtmp_header.body_length,
    (AMFState#rtmp_amf_state.header)#rtmp_header.type == ?RTMP_TYPE_INVOKE;
    (AMFState#rtmp_amf_state.header)#rtmp_header.type == ?RTMP_TYPE_NOTIFY ->
  Body = amf_decode(list_to_binary( lists:reverse( AMFState#rtmp_amf_state.chunks))),
  ProtocolState#rtmp_protocol_state.client_pid ! {rtmp, AMFState#rtmp_amf_state.header, Body},
  ProtocolState#rtmp_protocol_state{
    amf_states = store_amf(ProtocolState, reset(AMFState))
  };
process(ProtocolState, AMFState)
  when
    AMFState#rtmp_amf_state.chunks_length == (AMFState#rtmp_amf_state.header)#rtmp_header.body_length ->
  BodyBytes = list_to_binary( lists:reverse( AMFState#rtmp_amf_state.chunks)),
  ProtocolState#rtmp_protocol_state.client_pid ! {rtmp, AMFState#rtmp_amf_state.header, BodyBytes},
   BytesRead = ProtocolState#rtmp_protocol_state.bytes_read + (AMFState#rtmp_amf_state.header)#rtmp_header.body_length,
  send(ProtocolState#rtmp_protocol_state.socket, rtmp_bytesread(BytesRead)),
  ProtocolState#rtmp_protocol_state{
    amf_states = store_amf(ProtocolState, reset(AMFState)),
    bytes_read = BytesRead
  };
process(ProtocolState, AMFState) ->
  ProtocolState#rtmp_protocol_state{
    amf_states = store_amf(ProtocolState, AMFState)
  }.

store_amf(ProtocolState, AMFState) ->
  AMFIndex = (AMFState#rtmp_amf_state.header)#rtmp_header.amf_index,
  dict:store(AMFIndex, AMFState, ProtocolState#rtmp_protocol_state.amf_states).

reset(AMFState) ->
  AMFState#rtmp_amf_state{chunks = [], chunks_length = 0}.

handshake(Socket) ->
  ?D({handshake}),
  send(Socket, <<?HS_HEADER:8>>),
  send(Socket, handshake_bytes()),
  case get(Socket, 1) of
    <<?HS_HEADER:8>> ->
      ?D({handshake,header_found}),
      ServerBytes = get(Socket, ?HS_BODY_LEN),
      _ClientBytes = get(Socket, ?HS_BODY_LEN),
      send(Socket, ServerBytes),
      ?D({handshake,ok}),
      ok;
    _ ->
      ?D({rtmp,handshake,header_not_found}),
      exit(rtmp_header_not_found)
  end.

handshake_bytes() ->
  {Uptime, _} = statistics(runtime),
  <<Uptime:32/big, 0:32, (random_bytes(Uptime rem 256))/binary>>.
random_bytes(Seed) ->
  random_bytes(Seed, []).
random_bytes(_Seed, Bytes) when length(Bytes) =:= 764 ->
  list_to_binary(lists:flatten(lists:reverse(Bytes)));
random_bytes(Seed, Bytes) ->
  P = (12111221 * Seed + 1) rem 256,
  random_bytes(P, [[P,0]|Bytes]).

wait_result() ->
  receive
    {rtmp, _Header, ["_result"|_Rest] = Result} ->
      {ok, Result};
    {rtmp, _Header, ["_error"|_Rest] = Result} ->
      {error, Result};
    {rtmp, Header, Body} ->
      ?D({rtmp,still_waiting,but_not_for,{rtmp,Header,Body}}),
      wait_result()
  end.

rtmp_ping_pong(Time) ->
  ems_rtmp:encode( #channel{id=2, timestamp = 0, type = ?RTMP_TYPE_PING, stream = 0}, <<?PING_PONG_SERVER:16, Time:32>>).

rtmp_ping_clientbuffer(Object, BufferTime) ->
  ems_rtmp:encode( #channel{id=2, timestamp = 0, type = ?RTMP_TYPE_PING, stream = 0}, <<?PING_CLIENT_BUFFER:16, Object:32, BufferTime:32>>).

rtmp_serverbw() ->
  ems_rtmp:encode( #channel{id=2, timestamp = 0, type = ?RTMP_TYPE_BW_SERVER, stream = 0}, <<16#1312d0:32>>).

rtmp_bytesread(BytesRead) ->
  ems_rtmp:encode( #channel{id=2, timestamp = 0, type = ?RTMP_TYPE_BYTES_READ, stream = 0}, <<BytesRead:32>>).

amf_connect({stream, Host, {_Type, Id}, Auth}) ->
  App = lists:flatten(io_lib:format(
    "rtmp://~s:1935/ondemand?_fcs_vhost=~s&auth=~s&aifp=v001&slist=~s",
    [Host, Host, Auth, Id]
  )),
  #amf{id = 1, type = invoke, command = "connect", args = [
    {app, App},
    {flastVer,"MAC.9,0,47,0"},
    {swfUrl,"http://www.bbc.co.uk/emp/player.swf?revision=3887"},
    {tcUrl, App},
    {fpad,false},
    {capabilities,15},
    {audioCodecs,615},
    {videoCodecs,76},
    {pageUrl,"http://www.bbc.co.uk/iplayer/episode/b00ccd73"}]
  }.

amf_create_stream(StreamId) ->
  #amf{type = notify, command = "createStream", args = {list, [StreamId, null]}}.

amf_play({stream, _Host, {Type, Id}, Auth}) ->
  AMF_CALL_NETSTREAM_PLAY = 0,
  #amf{type = notify, command = "play", args = {list, [AMF_CALL_NETSTREAM_PLAY, null, 
    lists:flatten(
      io_lib:format("~s:~s?auth=~s&aifp=v001",
      [Type,Id,Auth])
    )
  ]}}.

amf_decode(AMF) -> amf_decode(AMF, []).
amf_decode(<<>>, Accum) -> Accum;
amf_decode(AMF, Accum) -> 
  {Object, Rest} = case ems_amf:parse(AMF) of
    {O, R} ->
      {O, R};
    {_Type, O, R} ->
      {O, R}
  end,
  amf_decode(Rest, Accum ++ [Object]).

flv_decode_audio_tag(Tag) ->
  <<_:1, Format:3, Rate:2, Bits:1, Channels:1>> = <<Tag:8>>,
  {flv_audio_format(Format), flv_audio_rate(Rate), flv_audio_bits(Bits), flv_audio_channels(Channels)}.

flv_audio_format(0) -> uncompressed;
flv_audio_format(1) -> adpcm;
flv_audio_format(2) -> mp3;
flv_audio_format(5) -> nellymoser_8khz;
flv_audio_format(6) -> nellymoser.

flv_audio_rate(0) -> 5.5;
flv_audio_rate(1) -> 11;
flv_audio_rate(2) -> 22;
flv_audio_rate(3) -> 44.

flv_audio_bits(0) -> 8;
flv_audio_bits(1) -> 16.

flv_audio_channels(0) -> mono;
flv_audio_channels(1) -> stereo.

rtmp_decode_metaheader(<<0:2,Index:6>>) -> {12, Index};
rtmp_decode_metaheader(<<1:2,Index:6>>) -> {8, Index};
rtmp_decode_metaheader(<<2:2,Index:6>>) -> {4, Index};
rtmp_decode_metaheader(<<3:2,Index:6>>) -> {1, Index}.

rtmp_decode_header(<<_MetaHeader:8, _AV:24, Length:24, Type:8, Stream:32/little>>, CurrentHeader) ->
  CurrentHeader#rtmp_header{body_length = Length, type = Type, stream = Stream};
rtmp_decode_header(<<_MetaHeader:8, _AV:24, Length:24, Type:8>>, CurrentHeader) ->
  CurrentHeader#rtmp_header{body_length = Length, type = Type};
rtmp_decode_header(<<_MetaHeader:8, _AV:24>>, CurrentHeader) ->
  CurrentHeader;
rtmp_decode_header(<<_MetaHeader:8>>, CurrentHeader) ->
  CurrentHeader.

rtmp_get_chunk(_Socket, 0, _ChunkSize) ->
  <<>>;
rtmp_get_chunk(Socket, RemainingLength, ChunkSize) when RemainingLength =< ChunkSize ->
  get(Socket, RemainingLength);
rtmp_get_chunk(Socket, RemainingLength, ChunkSize) when RemainingLength > ChunkSize ->
  get(Socket, ChunkSize).

send(Socket, Bytes) ->
  ok = gen_tcp:send(Socket, Bytes).

get(Socket, Size) ->
  case gen_tcp:recv(Socket, Size) of
    {ok, Bytes} ->
      Bytes;
    {error, Reason} ->
      exit(Reason)
  end.
